
from configure import *
from PIL import Image,ImageDraw
from resize import *
from torchvision.transforms import transforms
from yolo_net import *
from torch import  nn 
import numpy as np
tf = transforms.Compose([transforms.ToTensor()])
device = torch.device('cuda' if torch.cuda.is_available else 'cpu')
cls_num_pre = {0:'have_mask', 1:'no_mask'}

def NMS(dets, thresh):
    x1 = dets[:, 1]-0.5*dets[:, 3]      #[204,257,280,255]
    y1 = dets[:, 2]-0.5*dets[:, 4]        #[102,118,135,118]
    x2 = dets[:, 1]+0.5*dets[:, 3]         #[358,380,400,360]
    y2 = dets[:, 2]+0.5*dets[:, 4]          #[250,250,250,235]
    scores = dets[:, 0]     #[0.5,0.7,0.6,0.7]
    
    areas = (x2 - x1 + 1) * (y2 - y1 + 1) # 每个boundingbox的面积
    order = scores.argsort()[::-1] # boundingbox的置信度排序[3,1,2,0]
    keep = [] # 用来保存最后留下来的boundingbox
    while order.size > 0:     
        i = order[0] # 置信度最高的boundingbox的index
        keep.append(i) # 添加本次置信度最高的boundingbox的index
        
        # 当前bbox和剩下bbox之间的交叉区域
        # 选择大于x1,y1和小于x2,y2的区域
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])
        
        # 当前bbox和其他剩下bbox之间交叉区域的面积
        w = np.maximum(0.0, xx2 - xx1 + 1)
        h = np.maximum(0.0, yy2 - yy1 + 1)
        inter = w * h
        
        # 交叉区域面积 / (bbox + 某区域面积 - 交叉区域面积)
        ovr = inter / (areas[i] + areas[order[1:]] - inter)
        
        #保留交集小于一定阈值的boundingbox
        inds = np.where(ovr <= thresh)[0]
      
        order = order[inds + 1]
        
    return keep

class Detection(nn.Module):
    def __init__(self):
        super(Detection,self).__init__()
        self.weight = 'E:/python file/yolo_advance/weight/yolo_net.pt'
        # self.device = torch.device('cpu')
        self.model = Yolo_all()
        self.model.load_state_dict(torch.load(self.weight))
        self.model.eval()
    def forward(self,input,thresh,anchor_set,case):
        output13,output26,output52 = self.model(input)
        index13,box13 = self.get_index_box(output13,thresh)
        true_box13 = self.true_position(index13,box13,anchor_set[13],32.0,case)

        index26,box26 = self.get_index_box(output26,thresh)
        true_box26 = self.true_position(index26,box26,anchor_set[26],16.0,case)

        index52,box52 = self.get_index_box(output52,thresh)
        true_box52 = self.true_position(index52,box52,anchor_set[52],8.0,case)
        return torch.cat([true_box13,true_box26,true_box52],dim=0)
        
    def get_index_box(self,out,thresh):

        out = out.permute(0,2,3,1)
        out = out.reshape(out.size(0),out.size(1),out.size(2),3,-1)
        mask = torch.sigmoid(out[...,0])>thresh

        index = mask.nonzero()

        box = out[mask]
        return index,box
    def true_position(self,index,box,anchor_,k,case):
        anchor_ =  torch.Tensor(anchor_) 

        a = index[:,3]
        cy = (index[:,1].float() + torch.sigmoid(box[:,1]).float())*k/case
        cx = (index[:,2].float()+ torch.sigmoid(box[:,2]).float())*k/case

        w =anchor_[a,0]*torch.exp(box[:,3])/case
        h =anchor_[a,1]*torch.exp(box[:,4])/case

        p = box[:,0]
        class_num = torch.sigmoid(box[:,5:])
        name = torch.argmax(class_num,1)

        return torch.stack([p,cx,cy,w,h,name],dim=1)


if __name__ =="__main__":
    detect = Detection()

    number = '0002'
    local = 'image_nomask'
    image_date = Image.open(f'E:/python file/yolo_advance/JPEGImages/{local}/{number}.jpg')
    
    image = Yolo_resize(f'E:/python file/yolo_advance/JPEGImages/{local}/{number}.jpg')
    temp  = max(image_date.size)
    case = 416/temp

    image = tf(image)
    image = torch.unsqueeze(image,dim=0)

    result_ =  detect(image,0.5,anchor,case)
    real = NMS(result_[:,0:5].detach().numpy(),0.05)
    name = result_[real,5].detach().numpy()
    draw = ImageDraw.Draw(image_date)
    for i,box in enumerate(real):
        x1,y1,x2,y2 = result_[box][1]-0.5*result_[box][3], result_[box][2]-0.5*result_[box][4], result_[box][1]+0.5*result_[box][3], result_[box][2]+0.5*result_[box][4]
        x1 = x1 if x1>0 else 0
        y1 = y1 if y1>0 else 0
        x2 = x2 if x2<image_date.size[0] else image_date.size[0]
        y2 = y2 if y2<image_date.size[1] else image_date.size[1]
        draw.rectangle((x1,y1,x2,y2),outline='red',width= 3)
        draw.text( (x1,y1),cls_num_pre[name[i]])
        print(cls_num_pre[name[i]])
    image_date.show()